#include "WinApi.h"
#include <windows.h>

#define PW_32( n )*(DWORD *)(n)
#define PW_16( n )*(WORD *)(n)

namespace WinApi {
	constexpr DWORD HASH_NTDLL = ELFNoCaseHash("ntdll.dll");
	constexpr DWORD HASH_KERNEL32 = ELFNoCaseHash("KERNEL32.dll");
	constexpr DWORD HASH_KERNELBASE = ELFNoCaseHash("KERNELBASE.dll");
	//-------Add Other Dll
	/*
	constexpr DWORD HASH_USER32 = ELFNoCaseHash("USER32.dll");
	*/

	constexpr DWORD HASH_LoadLibraryW = CalcHash("LoadLibraryW");
	constexpr DWORD HASH_LoadLibraryA = CalcHash("LoadLibraryA");
	constexpr DWORD HASH_GetProcAddress = CalcHash("GetProcAddress");

	constexpr DWORD HASH_RtlInitUnicodeString = CalcHash("RtlInitUnicodeString");
	constexpr DWORD HASH_LdrLoadDll = CalcHash("LdrLoadDll");


	typedef struct _UNICODE_STR
	{
		USHORT Length;
		USHORT MaximumLength;
		PWSTR pBuffer;
	} UNICODE_STR, *PUNICODE_STR;

	typedef struct _PEB_LDR_DATA
	{
		DWORD dwLength;
		DWORD dwInitialized;
		LPVOID lpSsHandle;
		LIST_ENTRY InLoadOrderModuleList;
		LIST_ENTRY InMemoryOrderModuleList;
		LIST_ENTRY InInitializationOrderModuleList;
		LPVOID lpEntryInProgress;
	} PEB_LDR_DATA, *PPEB_LDR_DATA;

	typedef struct _RTL_USER_PROCESS_PARAMETERS {
		BYTE			Reserved1[16];
		PVOID			Reserved2[10];
		UNICODE_STR		ImagePathName;
		UNICODE_STR		CommandLine;
	} RTL_USER_PROCESS_PARAMETERS, *PRTL_USER_PROCESS_PARAMETERS;

	typedef struct __PEB
	{
		BYTE Reserved1[2];
		BYTE BeingDebugged;
#ifdef _WIN64
		BYTE Reserved2[21];
#else
		BYTE Reserved2[9];
#endif
		PPEB_LDR_DATA pLdr;
		PRTL_USER_PROCESS_PARAMETERS ProcessParameters;
		BYTE Reserved3[520];
		BYTE Reserved4[136];
		ULONG SessionId;
	} _PEB, *_PPEB;

	typedef struct _LDR_DATA_TABLE_ENTRY
	{
		LIST_ENTRY InMemoryOrderModuleList;
		LIST_ENTRY InInitializationOrderModuleList;
		PVOID DllBase;
		PVOID EntryPoint;
		ULONG SizeOfImage;
		UNICODE_STR FullDllName;
		UNICODE_STR BaseDllName;
		ULONG Flags;
		SHORT LoadCount;
		SHORT TlsIndex;
		LIST_ENTRY HashTableEntry;
		ULONG TimeDateStamp;
	} LDR_DATA_TABLE_ENTRY, *PLDR_DATA_TABLE_ENTRY;

	HMODULE GetModuleHandleFromPeb(CONST DWORD& dwHash)
	{

#ifdef _WIN64
		DWORD64 dwPeb = __readgsqword(0x60);
#else
		DWORD dwPeb = __readfsdword(0x30);
#endif

		PLDR_DATA_TABLE_ENTRY pDataTable = (PLDR_DATA_TABLE_ENTRY)((PPEB_LDR_DATA)((_PPEB)dwPeb)->pLdr)->InMemoryOrderModuleList.Flink;//uiBaseAddress->PLDR_DATA_TABLE_ENTRY

		while (pDataTable)
		{
			LPTSTR lpDllName = (LPTSTR)pDataTable->BaseDllName.pBuffer;
			if (nullptr == lpDllName) {
				return nullptr;
			}

			USHORT nLen = pDataTable->BaseDllName.Length;

			DWORD dwValue = ELFNoCaseHash(lpDllName, nLen);

			if ((DWORD)dwValue == dwHash)
			{
				return (HMODULE)pDataTable->DllBase;
			}

			pDataTable = (PLDR_DATA_TABLE_ENTRY)pDataTable->InMemoryOrderModuleList.Flink;
		}

		return nullptr;
	}

	HMODULE GetModuleHandleFromDllNum(CONST DWORD& dwDllNum)
	{
		HMODULE hModule = GetModuleHandleFromPeb(dwDllNum);
		if (hModule) {
			return hModule;
		}

		switch (dwDllNum) {
		case HASH_KERNEL32:
		{
			HMODULE hNtdll = GetModuleHandleFromPeb(HASH_NTDLL);
			if (!hNtdll) {
				return nullptr;
			}
			NTSTATUS NTAPI LdrLoadDll(
					IN PWSTR DllPath OPTIONAL,
					IN PULONG DllCharacteristics OPTIONAL,
					IN PUNICODE_STR DllName,
					OUT PVOID *DllHandle );
			NTSTATUS NTAPI RtlInitUnicodeString(PUNICODE_STR, PCWSTR);

			auto pRtlInitUnicodeString = (decltype(&RtlInitUnicodeString))GetExportFunAddress(hNtdll, HASH_RtlInitUnicodeString);
			auto pLdrLoadDll = (decltype(&LdrLoadDll))GetExportFunAddress(hNtdll, HASH_LdrLoadDll);
			UNICODE_STR name;
			//Kernel32.dll
			WCHAR strKernel32[] = { L'K', L'e', L'r', L'n', L'e', L'l', L'3', L'2', L'.', L'd', L'l', L'l', L'\0', };

			pRtlInitUnicodeString(&name, strKernel32);
			HMODULE hKernel32;
			pLdrLoadDll(nullptr, 0, &name, (PHANDLE)&hKernel32);

			return hKernel32;
		}
		case HASH_KERNELBASE:
		{
			//Kernelbase.dll
			WCHAR strKernelBase[] = { L'K', L'e', L'r', L'n', L'e', L'l', L'b', L'a', L's', L'e', L'.', L'd', L'l', L'l', L'\0', };
			return API(KERNEL32, LoadLibraryW)(strKernelBase);
		}
		break;
		//-------Add Other Dll
		/* case HASH_USER32:
		{
			
			WCHAR strUser32[] = ...;
			return API(KERNEL32, LoadLibraryW)(strUser32);
		}
		break;
		*/

		default:
			break;
		}

		return nullptr;
	}

	PVOID GetExportFunAddress(CONST VOID* lpReflectiveDllBuffer, CONST DWORD& dwHash)
	{
		DWORD dwCounter = 0;
#ifdef _WIN64
		DWORD dwCompiledArch = 2;
#else
		DWORD dwCompiledArch = 1;
#endif

		UINT_PTR pBase = (UINT_PTR)lpReflectiveDllBuffer;

		UINT_PTR pExportDir = pBase + ((PIMAGE_DOS_HEADER)pBase)->e_lfanew;

		if (((PIMAGE_NT_HEADERS)pExportDir)->OptionalHeader.Magic == 0x010B)
		{
			if (dwCompiledArch != 1)
				return nullptr;
		}
		else if (((PIMAGE_NT_HEADERS)pExportDir)->OptionalHeader.Magic == 0x020B)
		{
			if (dwCompiledArch != 2)
				return nullptr;
		}
		else
		{
			return nullptr;
		}

		UINT_PTR pNameArray = (UINT_PTR)&((PIMAGE_NT_HEADERS)pExportDir)->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];

		pExportDir = pBase + ((PIMAGE_DATA_DIRECTORY)pNameArray)->VirtualAddress;
		DWORD dwSize = ((PIMAGE_DATA_DIRECTORY)pNameArray)->Size;

		pNameArray = pBase + ((PIMAGE_EXPORT_DIRECTORY)pExportDir)->AddressOfNames;

		UINT_PTR pAddressArray = pBase + ((PIMAGE_EXPORT_DIRECTORY)pExportDir)->AddressOfFunctions;

		UINT_PTR pNameOrdinals = pBase + ((PIMAGE_EXPORT_DIRECTORY)pExportDir)->AddressOfNameOrdinals;

		dwCounter = ((PIMAGE_EXPORT_DIRECTORY)pExportDir)->NumberOfNames;

		while (dwCounter--)
		{
			LPCSTR cpExportedFunctionName = (LPCSTR)(pBase + PW_32(pNameArray));

			if (dwHash == CalcHash(cpExportedFunctionName))
			{
				pAddressArray = pBase + ((PIMAGE_EXPORT_DIRECTORY)pExportDir)->AddressOfFunctions;

				pAddressArray += (PW_16(pNameOrdinals) * sizeof(DWORD));
				LPBYTE lpAddr = (LPBYTE)lpReflectiveDllBuffer + PW_32(pAddressArray);

				if (lpAddr > (LPBYTE)pExportDir
					&& lpAddr < (LPBYTE)pExportDir + dwSize)
				{
					char* s = (char*)lpAddr;
					char nameDll[64];
					int i = 0;
					while (s[i] != '.' && i < 60)
					{
						nameDll[i] = s[i];
						i++;
					}
					if (i >= 60) {
						return nullptr;
					}
					s += i + 1;
					nameDll[i++] = '.';
					nameDll[i++] = 'd';
					nameDll[i++] = 'l';
					nameDll[i++] = 'l';
					nameDll[i] = 0;
					int num = 0;
					if (*s == '#')
					{
						while (*++s) {
							num = num * 10 + *s - '0';
						}
						s = (char*)&num;
					}

					HMODULE hDll = API(KERNEL32, LoadLibraryA)(nameDll);
					if(hDll){
						return API(KERNEL32, GetProcAddress)(hDll, s);
					}
					return nullptr;
				}
				else {
					return (PVOID)lpAddr;
				}
			}

			pNameArray += sizeof(DWORD);

			pNameOrdinals += sizeof(WORD);
		}

		return nullptr;
	}


	PVOID GetApiAddr(DWORD dwModuleNum, DWORD dwFunHash)
	{
		HMODULE hModule = GetModuleHandleFromDllNum(dwModuleNum);

		if (hModule) {
			return (PVOID)GetExportFunAddress((PBYTE)hModule, dwFunHash);
		}

		return nullptr;
	}
};
